#include <CImg.h>
#include <Primative.h>
#include <Scene.h>
#include <memory>
#include <clUtil.h>
#include <vector>
#include <sys/time.h>
#include <exception>

using namespace cimg_library;
using namespace std;
using namespace clUtil;

const unsigned int height = 2160;//1080;
const unsigned int width = 3840;//1920;
  
float image[4 * height * width];

double getTime()
{
  struct timeval tv;

  gettimeofday(&tv, NULL);

  return tv.tv_sec + tv.tv_usec / 1e6;
}

int main(int argc, char** argv)
{
  try
  {
    CImg<float> picture(width, height, 1, 3, 0.0f);

    const char* files[] = {"kernel.cl"};

    Device::InitializeDevices(files, 1);
    vector<unique_ptr<Buffer>> deviceBuffers;

    //Allocate buffers on each device
    for(size_t curDevice = 0; 
        curDevice < Device::GetDevices().size() * 2; 
        curDevice++)
    {
      Device::SetCurrentDevice(curDevice / 2);

      unique_ptr<Buffer> curBuffer(new Buffer(height * 
                                              width * 
                                              sizeof(float) * 
                                              4));
      deviceBuffers.push_back(move(curBuffer));
    }

    Device::Finish();


    //Raytrace that bitch using all your might!
    ParallelFor(0, 1, height - 1, [&](size_t start, size_t end)
    {
      size_t curDevice = Device::GetCurrentDeviceNum();
      unsigned int count = end - start + 1;
      unsigned int offset = (unsigned int) start;

      clUtilEnqueueKernel("raytracer", 
                          clUtilGrid(width, 16, count, 16), 
                          *deviceBuffers[2 * curDevice],
                          width, 
                          height,
                          count,
                          offset);
#if 0
      deviceBuffers[curDevice]->get(&image[4 * start * width], 
                                    4 * sizeof(float) * count * width);
#endif
    }, StaticScheduler(15));
   
    size_t count = 30;

    vector<double> times;

    for(size_t i = 0; i < count; i++)
    {
      double time1 = getTime();

      //Raytrace that bitch using all your might!
      ParallelFor(0, 1, height - 1, [&](size_t start, size_t end)
      {
        size_t curDevice = Device::GetCurrentDeviceNum();
        unsigned int count = end - start + 1;
        unsigned int offset = (unsigned int) start;
   
        size_t localSize;
        size_t curQueue = 0;

        if(Device::GetDevices()[curDevice].getDeviceInfo().Type ==
           CL_DEVICE_TYPE_CPU)
        {
          localSize = 16;
        }
        else
        {
          localSize = 16;
        }

        for(size_t curRow = 0; curRow < count; curRow += 64)
        {
          Device::GetCurrentDevice().setCommandQueue(curQueue);

          unsigned int innerRowCount = count - curRow > 64 ? 
            64 : count - curRow;
          unsigned int innerOffset = offset + curRow;    

          clUtilEnqueueKernel("raytracer", 
                              clUtilGrid(width, localSize, count, localSize), 
                              *deviceBuffers[2 * curDevice + curQueue],
                              width, 
                              height,
                              innerRowCount,
                              innerOffset);
#if 1
          deviceBuffers[2 * curDevice + curQueue]->
            get(&image[4 * innerOffset * width], 
                4 * sizeof(float) * innerRowCount * width);
#endif     
          curQueue = curQueue == 0 ? 1 : 0;
        }
      },
      EGSScheduler()); 
      //PINAScheduler("RaytraceLoop", 2));
      //StaticScheduler(6));

      double time2 = getTime();

      times.push_back(time2 - time1);
    }

    double minTime = 1e20;
    double maxTime = 0.0;
    double totalTime = 0.0f;

    for(size_t i = 0; i < times.size(); i++)
    {
      totalTime += times[i];
      minTime = times[i] < minTime ? times[i] : minTime;
      maxTime = times[i] > maxTime ? times[i] : maxTime;
    }

    cout << "Rays/sec: " << count * height * width / totalTime << endl;
    cout << "Max time: " << maxTime << endl;
    cout << "Min time: " << minTime << endl;
    cout << "Avg time: " << totalTime / count << endl;

    //Pack the buffer into a bitmap
    for(size_t i = 0; i < width; i++)
    {
      for(size_t j = 0; j < height; j++)
      {
        picture(i, j, 0) = 255.0f * image[4 * (i + j * width) + 0];
        picture(i, j, 1) = 255.0f * image[4 * (i + j * width) + 1];
        picture(i, j, 2) = 255.0f * image[4 * (i + j * width) + 2];
      }
    }

    picture.save_bmp("out.bmp");
  }
  catch(clUtilException& e)
  {
    cout << e.what();
    return 1;
  }

  return 0;
}

